import os
os.environ['CUDA_VISIBLE_DEVICES']='4'
import util_for_huggingface
from transformers.image_utils import load_image


import numpy as np
import torch

from MODEL_CKP import DINO3

def get_dino3():
    from transformers import pipeline
    feature_extractor = pipeline(
        model=DINO3,
        task="image-feature-extraction", 
    )
    return feature_extractor
def get_dino3_processor():
    from transformers import AutoImageProcessor, AutoModel
    pretrained_model_name = DINO3
    processor = AutoImageProcessor.from_pretrained(pretrained_model_name)
    model = AutoModel.from_pretrained(
        pretrained_model_name, 
        device_map="cuda", 
    )
    return processor,model

def get_feature_by_dino3(url='',img=None,feature_extractor = None):
    # url = "tmp.jpg"
    image = load_image(url) if img is None else img


    features = feature_extractor(image) # 1 201 4096
    return np.array( features )

def get_pooled_feature_by_dino3(url='',img=None,processor=None,model=None):
    image = load_image(url) if img is None else img
   
    inputs = processor(images=image, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        outputs = model(**inputs)

    pooled_output = outputs.pooler_output
    
    return pooled_output  # torch.Size([1, 4096])